import os
import numpy as np
import re


def weighted_average(existing_flow, new_flow, existing_weight, new_weight):
    total_weight = existing_weight + new_weight
    weighted_flow = (existing_flow * existing_weight + new_flow * new_weight) / total_weight
    return weighted_flow, total_weight


def split_and_save_flow(flow_dir, output_dir, tile_size=256):
    os.makedirs(output_dir, exist_ok=True)

    flow_files = [f for f in os.listdir(flow_dir) if f.endswith('.npy')]
    weight_dict = {}  # 存储每个小块的累计权重

    for flow_file in flow_files:
        flow_path = os.path.join(flow_dir, flow_file)
        flow = np.load(flow_path)

        h, w, _ = flow.shape
        # assert h == w == 512, "光流场的尺寸必须为512x512"

        # 使用正则表达式解析文件名中的坐标
        match = re.search(r'tile_\((\d+)_(\d+),(\d+)_(\d+)\)\.npy', flow_file)

        if not match:
            raise ValueError(f"文件名格式不正确: {flow_file}")

        top_left_row, top_left_col, bottom_right_row, bottom_right_col = map(int, match.groups())

        # 计算中间的坐标
        mid_row = top_left_row + 1
        mid_col = top_left_col + 1

        # 划分光流场为四个256x256的小块
        flows = {
            f'tile_({top_left_row}_{top_left_col})': flow[:tile_size, :tile_size],
            f'tile_({top_left_row}_{mid_col})': flow[:tile_size, tile_size:],
            f'tile_({mid_row}_{top_left_col})': flow[tile_size:, :tile_size],
            f'tile_({mid_row}_{mid_col})': flow[tile_size:, tile_size:]
        }

        for key, sub_flow in flows.items():
            sub_flow_path = os.path.join(output_dir, f"{key}.npy")
            if os.path.exists(sub_flow_path):
                existing_sub_flow = np.load(sub_flow_path)
                existing_weight = weight_dict.get(key, 1)
                new_weight = 1 / (existing_weight + 1)
                updated_flow, updated_weight = weighted_average(existing_sub_flow, sub_flow, existing_weight,
                                                                new_weight)
                weight_dict[key] = updated_weight
                np.save(sub_flow_path, updated_flow)
            else:
                weight_dict[key] = 1
                np.save(sub_flow_path, sub_flow)


# 示例用法
flow_dir = r'E:\wafer52\11867_32nm_fine_align_512_dvf'  # 替换为实际路径
output_dir = r'E:\wafer52\11867_32nm_fine_align_256_dvf_repair'  # 替换为实际路径

split_and_save_flow(flow_dir, output_dir, tile_size=256)

print(f"Processed flow fields saved to {output_dir}")
